#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/05 

from pathlib import Path
from collections import Counter
from typing import List

import matplotlib.pyplot as plt

from data import load_pdb_data, DATA_PATH


def seq_id_to_letter(seq:List[int], residue_names) -> str:
  return ''.join([residue_names[r] for r in seq])


if __name__ == '__main__':
  residue_names, sequences, locations = load_pdb_data()

  out_dp = Path(DATA_PATH) / 'pdb_stats'
  out_dp.mkdir(exist_ok=True)

  ''' plot_residue_freq '''
  residues = []
  for seq in sequences:
    residues.extend([r for r in seq])
  freq_residue = sorted([(c, r) for r, c in Counter(residues).items()], reverse=True)

  plt.clf()
  plt.plot([c for c, _ in freq_residue])
  plt.xticks(range(len(residue_names)), [r for c, r in freq_residue])
  fp = out_dp / 'residue_freq.png'
  print(f'>> save figure to {fp}')
  plt.savefig(fp, dpi=400)

  ''' list_uniq_chains '''
  freq_seq = sorted([(c, s) for s, c in Counter(sequences).items()], reverse=True)
  print('>> top 10 freq seqs:')
  for i in range(10):
    print(f'      ({freq_seq[i][0]}) {freq_seq[i][1]}')
  
  seq_counts  = [c for c, _ in freq_seq]
  seq_strings = [s for _, s in freq_seq]

  plt.clf()
  plt.yscale('log')
  plt.plot(seq_counts)
  fp = out_dp / 'sequence_uniq.png' 
  print(f'>> save figure to {fp}')
  plt.savefig(fp, dpi=400)

  print(f'>> found {len(seq_strings)} uniq seqs')

  fp = out_dp / 'sequence_uniq.txt' 
  print(f'>> save list to {fp}')
  with open(fp, 'w', encoding='utf-8') as fh:
    for s in seq_strings:
      fh.write(s)
      fh.write('\n')

  fp = out_dp / 'sequence_uniq_sorted.txt' 
  print(f'>> save sorted list to {fp}')
  with open(fp, 'w', encoding='utf-8') as fh:
    for s in sorted(seq_strings):
      fh.write(s)
      fh.write('\n')
